﻿//-------------------------------------
// WBFSSync - WBFSSync.exe
//
// Copyright 2009 Caian (ÔmΣga Frøst) <frost.omega@hotmail.com> :
//
// WBFSSync is Licensed under the terms of the Microsoft Reciprocal License (Ms-RL)
//
// DeviceDump.cs:
//
// Classe que engloba ferramentas de verificação de disco wbfs
//
//-------------------------------------

using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading;
using WBFSSync;
using System.IO;
using System.Net;

namespace DevConsole
{
    static class DeviceDump
    {
        const int wbfsMagic = (('W' << 24) | ('B' << 16) | ('F' << 8) | ('S'));
        const uint wiidiscMagic = 0x5D1C9EA3;

        const ushort discHeaderCopySize = 256;

        const byte wiiSectorSize_s = 15;
        const uint wiiSectorSize = 32768;
        const uint wiiSectorsPerDisc = 143432 * 2; //Suporte a Double Layer

        const int wbfsHeadMagicPos = 0; // Posição do "Magic" dentro do cabeçalho do sistema de arquivos
        const int wbfsHeadHdSecsPos = 4; // Posição do "Número de setores" dentro do cabeçalho do sistema de arquivos
        const int wbfsHeadHdSecSz = 8; // Posição do "Tamanho do setor" dentro do cabeçalho do sistema de arquivos
        const int wbfsHeadWbfsSecSz = 9; // Posição do "Tamanho do setor WBFS" dentro do cabeçalho do sistema de arquivos
        const int wbfsHeadDiscTable = 12; // Posição da "Tabela de disco" dentro do cabeçalho do sistema de arquivos

        public const int wiidiscMagicPos = 24;
        public const uint wbfsHeaderSize = 12;

        //-------------------------------
        static long step = 0;
        static long steps = 1;

        static List<String> log = new List<string>();
        static string file = "";

        //-------------------------------
        static void LogLine(String format, params Object[] args)
        {
            log.Add(String.Format(format, args));
            Draw();
        }

        //-------------------------------
        static void Log(String format, params Object[] args)
        {
            log[log.Count - 1] = log[log.Count - 1] + String.Format(format, args);
            Draw();
        }

        //-------------------------------
        public static void DumpCommand(String[] args)
        {
            if (args.Length != 2)
            {
                Console.WriteLine("dump DRIVE-LETTER");
                return;
            }

            Dump(args[1][0]);
        }

        //-------------------------------
        static void Dump(Char drive)
        {
            step = 0;
            steps = 1;

            string devfile = @"\\.\" + drive + ':';
            LogLine("Opening Drive '{0}'...", devfile);

            IIOContext device = IOManager.CreateIOContext("DEVVERIFY", devfile, System.IO.FileAccess.ReadWrite,
                FileShare.None, 0, FileMode.Open, EFileAttributes.NoBuffering);

            if (device.Result != 0)
            {
                LogLine("Failed to open device! Error {0}", device.Result);
                goto END;
            }

            if (device.Lock() != IORet.RET_IO_OK)
            {
                LogLine("Failed to open device! Error {0}", device.Result);
                goto END;
            }

            //

            DISK_GEOMETRY dg = device.GetDiskGeometry();
            PARTITION_INFORMATION pi = device.GetPartitionInformation();

            if (dg.BytesPerSector == 0)
            {
                LogLine("Device is empty!");
                goto END;
            }

            //

            LogLine("Loading partition Head...");

            uint PartitionOffsetLBA = 0;

            uint hdSectorSize = dg.BytesPerSector;
            uint hdTotalSectors = (uint)(pi.PartitionLength / hdSectorSize);

            byte hdSectorSize_s = bitshift((uint)hdSectorSize);
            byte hdTotalSectors_s = bitshift((uint)hdTotalSectors);

            uint wiiTotalSectors = (uint)((long)hdTotalSectors * hdSectorSize / wiiSectorSize);

            //Cabeçalho

            Byte[] Head = new Byte[hdSectorSize];

            //Lê o cabeçalho do disco
            if (device.Read(Head, PartitionOffsetLBA * hdSectorSize, (int)hdSectorSize) != IORet.RET_IO_OK)
            {
                LogLine("Failed to read device! Error {0}", device.Result);
                goto END;
            }
            else
            {
                Log("Ok");
            }

            uint discWbfsMagic = BitConverter.ToUInt32(Head, wbfsHeadMagicPos);
            uint discTotalSectors = ntohi(BitConverter.ToUInt32(Head, wbfsHeadHdSecsPos));
            byte discSectorSize_s = Head[wbfsHeadHdSecSz];

            if (discWbfsMagic != htoni(wbfsMagic))
            {
                LogLine("Incorrect partition magic, fixing...");
                Byte[] newmagic = BitConverter.GetBytes(htoni(wbfsMagic));
                newmagic.CopyTo(Head, wbfsHeadMagicPos);
                Log("Ok");
            }

            if (discTotalSectors != hdTotalSectors)
            {
                LogLine("Incorrect partition total sectors, fixing...");
                Byte[] newmagic = BitConverter.GetBytes(htoni(hdTotalSectors));
                newmagic.CopyTo(Head, wbfsHeadHdSecsPos);
                Log("Ok");
            }

            if (discSectorSize_s != hdSectorSize_s)
            {
                LogLine("Incorrect partition sector size, fixing...");
                Head[wbfsHeadHdSecSz] = hdSectorSize_s;
                Log("Ok");
            }

            //

            Byte size_s;
            for (size_s = 6; size_s < 11; size_s++)
            {
                if (wiiTotalSectors < ((1U << 16) * (1 << size_s)))
                {
                    break;
                }
            }

            byte wbfsSectorSize_s = (byte)(size_s + wiiSectorSize_s);
            Head[wbfsHeadWbfsSecSz] = wbfsSectorSize_s;

            //

            uint wbfsSectorSize = (uint)(1 << wbfsSectorSize_s);
            ushort wbfsTotalSectors = (ushort)((int)wiiTotalSectors >> (int)(wbfsSectorSize_s - wiiSectorSize_s));
            ushort wbfsSectorsPerDisc = (ushort)((int)wiiSectorsPerDisc >> (int)(wbfsSectorSize_s - wiiSectorSize_s));

            uint wiiSectorsPerWBFSSector = (uint)(1 << (int)(wbfsSectorSize_s - wiiSectorSize_s));

            ushort discInfoSize = (ushort)Align_LBA((int)(discHeaderCopySize + wbfsSectorsPerDisc * 2), hdSectorSize);
            uint discInfoSizeInSecs = (uint)(discInfoSize >> hdSectorSize_s);

            uint WLBAPositionLBA = (uint)(((int)wbfsSectorSize - (int)wbfsTotalSectors / 8) >> (int)hdSectorSize_s);
            int WLBASize = Align_LBA(wbfsTotalSectors / 8, hdSectorSize);

            ushort maxDiscs = Math.Min((ushort)((WLBAPositionLBA - 1) / (discInfoSize >> hdSectorSize_s)),
                (ushort)(hdSectorSize - wbfsHeaderSize));

            //

            LogLine("Loading WBFS Sector 0...");

            Byte[] WBFSHead = new Byte[wbfsSectorSize];

            if (device.Read(WBFSHead, PartitionOffsetLBA * hdSectorSize, (int)wbfsSectorSize) != IORet.RET_IO_OK)
            {
                LogLine("Failed to read device! Error {0}", device.Result);
                goto END;
            }
            else
            {
                Log("Ok");
            }

            //

            LogLine("Searching for disc heads...");
            steps = wbfsTotalSectors;
            step = 1;

            for (int i = 1; i < wbfsTotalSectors; i++)
            {
                step++;
                file = String.Format("{0}\\SECTORTABLE[{1}]", devfile, i);
                Draw();

                //

                long position = (long)wbfsSectorSize * i;
                Byte[] info = new byte[hdSectorSize];

                if (device.Read(info, position, (int)hdSectorSize) != IORet.RET_IO_OK)
                {
                    LogLine("Failed to read device! Error {0}", device.Result);
                    goto END;
                }

                //

                uint magic = _be32(info, 24);
                if (magic != 0x5D1C9EA3)
                {
                    //LogLine("Disc {0} is not a wii disc!", i);
                    continue;
                }

                //

                String code = Encoding.Default.GetString(info, 0, 6);
                String name = Encoding.Default.GetString(info, 32, 32).TrimEnd('\0');

                LogLine("Found possible disc head '{0} - {1}' in sector '{2}'", code, name, i);
            }

            //

            LogLine("Saving dump...");

            String coredumppath = Environment.GetFolderPath(Environment.SpecialFolder.DesktopDirectory) + "\\Dump\\";
            Directory.CreateDirectory(coredumppath);

            File.WriteAllBytes(coredumppath + "wbfssec0.dump", WBFSHead);
            File.WriteAllLines(coredumppath + "Log.txt", log.ToArray(), Encoding.Default);

            LogLine("Done.");

        END:

            Head = null;
            WBFSHead = null;

            device.Unlock();
            device.Close();

            Console.Write("Press Any Key");
            Console.ReadKey(true);
            Program.LoadMainScreen();
        }

        //-------------------------------
        static void Draw()
        {
            int i = 0;
            StringBuilder screen = new StringBuilder();

            //-------------------------------
            screen.Append('╔');
            for (i = 0; i < Console.WindowWidth - 2; i++) screen.Append('═');
            screen.Append('╗');

            screen.Append("║ Disk dump utility");
            for (i = 0; i < Console.WindowWidth - 20; i++) screen.Append(' ');
            screen.Append('║');

            screen.Append('╠');
            for (i = 0; i < Console.WindowWidth - 2; i++) screen.Append('═');
            screen.Append('╣');

            //-------------------------------
            int maxlog = Console.WindowHeight - 11;
            int logstart = log.Count - maxlog;

            for (i = Math.Max(logstart, 0); i < log.Count; i++)
            {
                screen.Append('║');

                string logline = log[i];
                if (logline.Length > Console.WindowWidth - 2) logline = logline.Substring(0, Console.WindowWidth - 2);
                screen.Append(logline);

                for (int j = 0; j < Console.WindowWidth - logline.Length - 2; j++) screen.Append(' ');
                screen.Append('║');
            }

            //-------------------------------
            while (i++ < maxlog)
            {
                screen.Append('║');
                for (int j = 0; j < Console.WindowWidth - 2; j++) screen.Append(' ');
                screen.Append('║');
            }

            //-------------------------------
            screen.Append('╠');
            for (i = 0; i < Console.WindowWidth - 2; i++) screen.Append('═');
            screen.Append('╣');

            screen.Append('║');
            for (int j = 0; j < Console.WindowWidth - 2; j++) screen.Append(' ');
            screen.Append('║');

            //-------------------------------
            screen.Append("║ ");

            string fileline = file;
            if (fileline.Length > Console.WindowWidth - 4) fileline = fileline.Substring(0, Console.WindowWidth - 4);
            screen.Append(fileline);

            for (int j = 0; j < Console.WindowWidth - fileline.Length - 4; j++) screen.Append(' ');
            screen.Append(" ║");

            //-------------------------------
            screen.Append("║ ");
            float progress = (float)(Console.WindowWidth - 4) * step / (float)steps;
            for (i = 0; i < (int)progress; i++) screen.Append('█');
            while (i++ < Console.WindowWidth - 4) screen.Append('░');
            screen.Append(" ║");

            screen.Append('║');
            for (int j = 0; j < Console.WindowWidth - 2; j++) screen.Append(' ');
            screen.Append('║');

            screen.Append('╚');
            for (int j = 0; j < Console.WindowWidth - 2; j++) screen.Append('═');
            screen.Append('╝');

            //-------------------------------
            Console.Clear();
            Console.ResetColor();
            Console.Write(screen.ToString());
        }

        //-------------------------------
        static byte bitshift(uint size)
        {
            byte r = 0;

            do
            {
                size >>= 1;
                r++;

            } while (size > 1);

            return r;
        }

        static int Align_LBA(int x, uint hdSectorSize)
        {
            return ((x) + (int)hdSectorSize - 1) & (~((int)hdSectorSize - 1));
        }

        static uint _be32(byte[] p, uint i)
        {
            return ((uint)(p[i]) << 24) | ((uint)(p[i + 1]) << 16) | ((uint)(p[i + 2]) << 8) | (uint)(p[i + 3]);
        }

        static bool IsSectorUsed(Byte[] wbfssector0, int pos, ushort sector)
        {
            int abssector = (sector - 1) / 32;
            uint v = ntohi(BitConverter.ToUInt32(wbfssector0, pos + 4 * abssector));
            return (v & (1 << ((sector - 1) % 32))) == 0;
        }

        //----------------- Network to Host
        static uint ntohi(uint i) { return (uint)IPAddress.NetworkToHostOrder((int)i); }
        static ushort ntohs(ushort i) { return (ushort)IPAddress.NetworkToHostOrder((short)i); }
        static ulong ntohl(ulong i) { return (ulong)IPAddress.NetworkToHostOrder((long)i); }

        //----------------- Host to Network
        static uint htoni(uint i) { return (uint)IPAddress.HostToNetworkOrder((int)i); }
        static ushort htons(ushort i) { return (ushort)IPAddress.HostToNetworkOrder((short)i); }
        static ulong htonl(ulong i) { return (ulong)IPAddress.HostToNetworkOrder((long)i); }
    }
}
